from turtle import forward
import numpy as np

from .module import Module
from autograd import Tensor
from autograd.nn.parameter import Parameter
from .. import functional as F

from typing import Callable, Optional


class _Loss(Module):
    reduction: str

    def __init__(self, reduction: str = 'mean') -> None:
        super(_Loss, self).__init__()
        self.reduction = reduction

class CrossEntropyLoss(_Loss):
    
    def __init__(self, reduction: str = 'mean') :
        super(CrossEntropyLoss, self).__init__()
        self.reduction = reduction

    def forward(self, input, target):
        return F.cross_entropy(input, target, self.reduction)

